from typing import List, Dict
from PIL import Image, ImageDraw, ImageFilter
import re

class MaskingService:
    """敏感信息打码服务"""
    
    def mask_text(self, text: str, detected_items: List[Dict]) -> str:
        """对文本进行打码处理"""
        masked_text = text
        
        # 按位置倒序排序，避免替换时位置偏移
        sorted_items = sorted(detected_items, key=lambda x: x.get('start', 0), reverse=True)
        
        for item in sorted_items:
            start = item.get('start', 0)
            end = item.get('end', 0)
            content = item.get('content', '')
            
            if start < end and start < len(masked_text):
                # 生成掩码
                mask = self._generate_mask(content, item.get('type', ''))
                masked_text = masked_text[:start] + mask + masked_text[end:]
        
        return masked_text
    
    def mask_image(self, image: Image.Image, detected_items: List[Dict]) -> Image.Image:
        """对图片进行打码处理"""
        masked_image = image.copy()
        draw = ImageDraw.Draw(masked_image)
        
        for item in detected_items:
            bbox = item.get('bbox')
            if bbox:
                # 计算边界框
                x_coords = [point[0] for point in bbox]
                y_coords = [point[1] for point in bbox]
                x_min, x_max = int(min(x_coords)), int(max(x_coords))
                y_min, y_max = int(min(y_coords)), int(max(y_coords))
                
                # 绘制黑色矩形（打码）
                draw.rectangle([x_min, y_min, x_max, y_max], fill='black')
                
                # 或者使用高斯模糊
                # region = masked_image.crop((x_min, y_min, x_max, y_max))
                # blurred = region.filter(ImageFilter.GaussianBlur(radius=10))
                # masked_image.paste(blurred, (x_min, y_min))
        
        return masked_image
    
    def _generate_mask(self, content: str, item_type: str) -> str:
        """生成掩码文本"""
        length = len(content)
        
        if item_type == 'id_card':
            # 身份证号：保留前3位和后4位
            if length >= 7:
                return content[:3] + '*' * (length - 7) + content[-4:]
            else:
                return '*' * length
        elif item_type == 'phone':
            # 手机号：保留前3位和后4位
            if length >= 7:
                return content[:3] + '*' * (length - 7) + content[-4:]
            else:
                return '*' * length
        elif item_type == 'bank_card':
            # 银行卡：保留后4位
            if length >= 4:
                return '*' * (length - 4) + content[-4:]
            else:
                return '*' * length
        else:
            # 其他类型：全部掩码
            return '*' * length

